// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::aggregate::test;
using namespace kumo::pollux::exec::test;

namespace kumo::pollux::aggregate::test {

namespace {

class CountDistinctTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    prestosql::registerInternalAggregateFunctions("");
    AggregationTestBase::SetUp();
  }
};

TEST_F(CountDistinctTest, global) {
  auto data = make_row_vector({
      make_flat_vector<int32_t>({1, 2, 3, 4, 5, 3, 4, 2, 6, 7}),
      make_flat_vector<StringView>(
          {"1", "2", "3", "4", "5", "3", "4", "2", "6", "7"}),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>(std::vector<int64_t>{7}),
  });

  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c0)"}, {expected});
  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c1)"}, {expected});

  // Null inputs.
  data = make_row_vector({
      make_nullable_flat_vector<int32_t>(
          {1, 2, std::nullopt, 4, 5, std::nullopt, 4, 2, 6, 7}),
      make_nullable_flat_vector<StringView>(
          {"1", "2", std::nullopt, "4", "5", std::nullopt, "4", "2", "6", "7"}),
  });

  expected = make_row_vector({
      make_flat_vector<int64_t>(std::vector<int64_t>{6}),
  });

  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c0)"}, {expected});
  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c1)"}, {expected});

  // All inputs are null.
  data = make_row_vector({
      make_all_null_flat_vector<int32_t>(1'000),
      make_all_null_flat_vector<StringView>(1'000),
  });

  expected = make_row_vector({
      make_flat_vector<int64_t>(std::vector<int64_t>{0}),
  });

  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c0)"}, {expected});
  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c1)"}, {expected});
}

TEST_F(CountDistinctTest, groupBy) {
  auto data = make_row_vector({
      make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
      make_flat_vector<int32_t>({1, 2, 3, 4, 5, 3, 4, 2, 6, 7}),
      make_flat_vector<StringView>(
          {"1", "2", "3", "4", "5", "3", "4", "2", "6", "7"}),
  });

  auto expected = make_row_vector({
      make_flat_vector<int16_t>({1, 2}),
      make_flat_vector<int64_t>({4, 4}),
  });

  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c1)"},
      {expected});
  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c2)"},
      {expected});

  // Null inputs.
  data = make_row_vector({
      make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
      make_nullable_flat_vector<int32_t>(
          {1,
           std::nullopt,
           3,
           std::nullopt,
           5,
           std::nullopt,
           3,
           std::nullopt,
           6,
           1}),
      make_nullable_flat_vector<StringView>(
          {"1",
           std::nullopt,
           "3",
           std::nullopt,
           "5",
           std::nullopt,
           "3",
           std::nullopt,
           "6",
           "1"}),
  });

  expected = make_row_vector({
      make_flat_vector<int16_t>({1, 2}),
      make_flat_vector<int64_t>({1, 3}),
  });

  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c1)"},
      {expected});
  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c2)"},
      {expected});

  // All inputs are null for a group.
  data = make_row_vector({
      make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
      make_nullable_flat_vector<int32_t>(
          {1,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           1}),
      make_nullable_flat_vector<StringView>(
          {"1",
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           "1"}),
  });

  expected = make_row_vector({
      make_flat_vector<int16_t>({1, 2}),
      make_flat_vector<int64_t>({1, 0}),
  });

  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c1)"},
      {expected});
  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c2)"},
      {expected});
}

TEST_F(CountDistinctTest, globalArray) {
  auto data = make_row_vector({
      make_array_vector<int32_t>({
          {1, 2, 3},
          {4, 5},
          {1, 2, 3},
          {3, 4, 2, 6, 7},
          {1, 2, 3},
          {4, 5},
      }),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>(std::vector<int64_t>{3}),
  });

  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c0)"}, {expected});
}

TEST_F(CountDistinctTest, groupByArray) {
  auto data = make_row_vector({
      make_flat_vector<int16_t>({1, 1, 1, 2, 2, 1}),
      make_array_vector<int32_t>({
          {1, 2, 3},
          {4, 5},
          {1, 2, 3},
          {3, 4, 2, 6, 7},
          {1, 2, 3},
          {4, 5},
      }),
  });

  auto expected = make_row_vector({
      make_flat_vector<int16_t>({1, 2}),
      make_flat_vector<int64_t>({2, 2}),
  });

  testAggregations(
      {data, data, data},
      {"c0"},
      {"\"$internal$count_distinct\"(c1)"},
      {expected});
}

TEST_F(CountDistinctTest, timestampWithTimeZone) {
  auto data = make_row_vector({
      // Keys for non-global aggregations.
      make_flat_vector<int16_t>({1, 1, 2, 1, 2, 1, 2, 1}),
      make_flat_vector<int64_t>(
          {pack(0, 0),
           pack(1, 0),
           pack(2, 0),
           pack(0, 1),
           pack(1, 1),
           pack(1, 2),
           pack(2, 2),
           pack(3, 3)},
          TIMESTAMP_WITH_TIME_ZONE()),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>(std::vector<int64_t>{4}),
  });

  testAggregations(
      {data}, {}, {"\"$internal$count_distinct\"(c1)"}, {expected});

  expected = make_row_vector({
      make_flat_vector<int16_t>({1, 2}),
      make_flat_vector<int64_t>({3, 2}),
  });

  testAggregations(
      {data}, {"c0"}, {"\"$internal$count_distinct\"(c1)"}, {expected});
}

} // namespace
} // namespace kumo::pollux::aggregate::test
